import numpy as np
import tensorflow as tf

class DataGenerator(tf.keras.utils.Sequence):
    """数据生成器 - 动态查找电影特征，避免爆内存
    支持 movie_idx 到 movieId 的反向映射。
    """

    def __init__(self, user_indices, movie_indices, user_features, movie_features_table, ratings, batch_size=1024, shuffle=True, inv_movie_id_map=None, **kwargs):
        """
        movie_features_table: DataFrame 或 dict，index/keys 为 movieId，值为特征向量
        movie_indices: 每条样本的 movie_idx（int，embedding输入）
        inv_movie_id_map: 可选，idx->movieId 的反向映射（如 np.array 或 list）
        """
        super().__init__(**kwargs)
        self.user_indices = user_indices
        self.movie_indices = movie_indices
        self.user_features = user_features
        self.movie_features_table = movie_features_table
        self.ratings = ratings
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.indexes = np.arange(len(self.ratings))
        # 反向映射，若未传入则假定 idx==movieId
        if inv_movie_id_map is not None:
            self.inv_movie_id_map = np.array(inv_movie_id_map)
        else:
            self.inv_movie_id_map = None
        self.on_epoch_end()

    def __len__(self):
        return int(np.ceil(len(self.ratings) / self.batch_size))

    def __getitem__(self, index):
        batch_indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]

        user_idx_batch = self.user_indices[batch_indexes]
        movie_idx_batch = self.movie_indices[batch_indexes]
        user_feat_batch = self.user_features[batch_indexes]
        # 用 movie_idx 映射回原始 movieId
        if self.inv_movie_id_map is not None:
            movie_id_batch = self.inv_movie_id_map[movie_idx_batch]
        else:
            movie_id_batch = movie_idx_batch  # 兼容 index==movieId 的情况
        # 动态查找每个 movieId 的特征
        if hasattr(self.movie_features_table, 'loc'):
            # DataFrame，只选取数值型特征
            numeric_cols = self.movie_features_table.select_dtypes(include=[np.number]).columns
            movie_feat_batch = self.movie_features_table.loc[movie_id_batch, numeric_cols].values.astype(np.float32)
        else:
            # dict 或 numpy
            movie_feat_batch = np.array([self.movie_features_table[mid] for mid in movie_id_batch], dtype=np.float32)
        ratings_batch = self.ratings[batch_indexes]

        return (user_idx_batch, movie_idx_batch, user_feat_batch, movie_feat_batch), ratings_batch

    def on_epoch_end(self):
        if self.shuffle:
            np.random.shuffle(self.indexes)

def create_tf_dataset(X, y, batch_size, shuffle=True, buffer_size=10000, prefetch_buffer=tf.data.AUTOTUNE):
    """
    使用 tf.data.Dataset 创建一个高效的数据管道。

    参数:
    - X (dict): 特征字典，键是特征名，值是 numpy 数组。
    - y (numpy.ndarray): 标签数组。
    - batch_size (int): 批次大小。
    - shuffle (bool): 是否打乱数据。
    - buffer_size (int): 用于打乱的缓冲区大小。
    - prefetch_buffer (int): 预取缓冲区大小，tf.data.AUTOTUNE 会自动调整。

    返回:
    - tf.data.Dataset: 配置好的 TensorFlow 数据集。
    """
    # 将输入字典和标签组合成一个数据集
    dataset = tf.data.Dataset.from_tensor_slices((X, y))

    if shuffle:
        # 打乱数据
        dataset = dataset.shuffle(buffer_size=buffer_size)

    # 将数据分批
    dataset = dataset.batch(batch_size)

    # 在后台预取数据，以加速训练
    dataset = dataset.prefetch(buffer_size=prefetch_buffer)

    return dataset